import random

import torch
import torch.nn as nn
import numpy as np

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")


class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return torch.nn.functional.softmax(x, dim=-1)


class Critic(nn.Module):
    def __init__(self, state_dim, hidden_dim=64):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class A2C:
    def __init__(self, state_dim, action_dim, action_prob_dim, lr_a, lr_c, epsilon, batch_size):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.action_prob_dim = action_prob_dim
        self.actor = Actor(state_dim, action_prob_dim).to(device)
        self.critic = Critic(state_dim).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr_a)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr_c)
        self.epsilon = epsilon
        self.batch_size = batch_size
        self.train_counter = 0
        self.buffer = []

    def select_action(self, require_print, state):
        p = np.random.rand()
        # 1. 贪心，随机抽样生成动作
        if p < self.epsilon:
            action_prob = np.random.random(size=self.action_dim)
            action_prob = np.column_stack((action_prob, 1 - action_prob))
            action = np.zeros(self.action_dim, dtype=np.int32)
            for i in range(self.action_dim):
                action[i] = np.random.choice([0, 1], p=action_prob[i])
            if require_print:
                avg_off_prob = np.mean(action_prob)
                off_ues_num = np.sum(action)
                print(f"\tChoose action by Greedy: distribution={avg_off_prob}, offloading_ues_num={off_ues_num}")
            return action
        # 2. 神经网络输出概率分布并随机抽样
        else:
            action_prob = self.actor(state)
            action_prob = (action_prob - action_prob.min()) / (action_prob.max() - action_prob.min())
            action = torch.bernoulli(action_prob)
            if require_print:
                avg_off_prob = torch.mean(action_prob).item()
                off_ues_num = torch.sum(action)
                print(f"\tChoose action by NN: distribution={avg_off_prob}, offloading_ues_num={off_ues_num}")
            return action.detach().numpy().astype(np.int32)

    def save_buffer(self, state, action, reward):
        self.buffer.append((state, action, reward))

    def learn(self):
        if len(self.buffer) < self.batch_size:
            return
        # 1. 随机抽取批次
        batch = random.sample(self.buffer, self.batch_size)
        states, actions, rewards = zip(*batch)
        # 2. 转换为张量
        states = torch.tensor(np.array(states), dtype=torch.float).to(device)
        actions = torch.tensor(np.array(actions), dtype=torch.int64).to(device)
        rewards = torch.tensor(np.array(rewards), dtype=torch.float).view(-1, 1).to(device)
        # 3. 状态动作价值评估网络
        states_value = self.critic(states)
        action_probs = self.actor(states)
        # 4. 计算 Loss 并更新网络
        # 4.1. 由于系统只有一个状态（初始位置分布），所以只需要当前状态越接近 rewards 越好
        critic_loss = nn.functional.mse_loss(states_value, rewards)
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        # 4.2. 同样的因为没有未来状态，只需要根据 Critic 的评估调整当前状态的预测动作概率分布
        actor_loss = -torch.mean((rewards - states_value).detach() * torch.log(action_probs))
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()
        # 5. 清理 Buffer
        del self.buffer[:]
